import mediapipe as mp
import cv2
import math

# 关键点连接列表
connection = [(0, 1), (1, 2), (2, 3), (3, 7), (0, 4), (4, 5),
              (5, 6), (6, 8), (9, 10), (11, 12), (11, 13),
              (13, 15), (15, 17), (15, 19), (15, 21), (17, 19),
              (12, 14), (14, 16), (16, 18), (16, 20), (16, 22),
              (18, 20), (11, 23), (12, 24), (23, 24), (23, 25),
              (24, 26), (25, 27), (26, 28), (27, 29), (28, 30),
              (29, 31), (30, 32), (27, 31), (28, 32), (0, 33)]

# 关节角度连接列表
angle = [(16, 14, 12),  # 0左侧肘关节
         (15, 13, 11),  # 1右侧肘关节
         (14, 12, 24),  # 2左侧腋窝
         (13, 11, 23),  # 3右侧腋窝
         (26, 24, 12),  # 4左侧髋关节
         (25, 23, 11),  # 5右侧髋关节
         (28, 26, 24),  # 6左侧膝关节
         (27, 25, 23),  # 7右侧膝关节
         (14, 12, 11),  # 8左侧肩膀
         (13, 11, 12)]  # 9右侧肩膀

# 初始化检测器
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(min_tracking_confidence=0.5, min_detection_confidence=0.8)


def detector(img):
    """
    关键点检测函数
    :param img: 输入图像
    :return: 关键点位置
    """
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)    # 将图像从BGR转换为RGB，方便检测器识别
    result = pose.process(img)   # 检测关键点
    # 如果没检测到，则返回空列表
    if result.pose_landmarks == None:
        return []

    # 通过图像的尺寸，计算出关键点所在位置
    h, w, c = img.shape
    landmarks = []
    for mark in result.pose_landmarks.landmark:
        x = int(w * mark.x)
        y = int(h * mark.y)
        if x < 0:
            x = 0
        if y < 0:
            y = 0
        landmarks.append((x, y))
    if landmarks[11] and landmarks[12]:
        landmarks.append((abs(int((landmarks[12][0] + landmarks[11][0]) / 2)), abs(int((landmarks[12][1] + landmarks[11][1]) / 2))))

    return landmarks


def draw_connections(img, landmarks):
    """
    关键点连接函数
    :param img: 输入图像
    :param landmarks: 关键点坐标列表
    :return: 返回连接图像
    """
    for c in connection:
        cv2.line(img, (landmarks[c[0]]), landmarks[c[1]], (255, 0, 0), thickness=2)
    return img


def draw_angle(landmarks):
    """
    计算关节角度函数
    :param landmarks: 关键点坐标列表
    :return: 关节角度列表
    """
    angles = []
    for i in angle:
        angles.append(__angle(landmarks[i[0]], landmarks[i[1]], landmarks[i[2]]))
    return angles


def __angle(p1, p2, p3):
    """
    计算角度(通过余弦定理)
    :param p1: 左侧点
    :param p2: 中间点
    :param p3: 右侧点
    :return: 中间夹角
    """
    a = ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5
    b = ((p2[0] - p3[0]) ** 2 + (p2[1] - p3[1]) ** 2) ** 0.5
    c = ((p1[0] - p3[0]) ** 2 + (p1[1] - p3[1]) ** 2) ** 0.5
    cosC = ((a ** 2) + (b ** 2) - (c ** 2)) / (2 * a * b + 1e-6)
    return math.acos(cosC) * 180 / 3.1415926   # 将弧度转化为角度
